package drill

import (
	"database/sql"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	_ "modernc.org/sqlite"
)

// DrillStateDbClient 钻井状态数据DAO工具类
type DrillStateDbClient struct {
	dataDir    string
	mu         sync.RWMutex
	partitions map[string]*sql.DB // 分区数据库缓存
}

// NewDrillStateDbClient 创建新的钻井状态数据库客户端
func NewDrillStateDbClient(dataDir string) (*DrillStateDbClient, error) {
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create data directory: %v", err)
	}

	client := &DrillStateDbClient{
		dataDir:    dataDir,
		partitions: make(map[string]*sql.DB),
	}

	return client, nil
}

// Close 关闭所有数据库连接
func (c *DrillStateDbClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	var errs []error
	for _, db := range c.partitions {
		if err := db.Close(); err != nil {
			errs = append(errs, err)
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("errors closing databases: %v", errs)
	}
	return nil
}

// getPartitionDB 获取指定月份的分区数据库
func (c *DrillStateDbClient) getPartitionDB(partitionKey string) (*sql.DB, error) {
	c.mu.RLock()
	if db, exists := c.partitions[partitionKey]; exists {
		c.mu.RUnlock()
		return db, nil
	}
	c.mu.RUnlock()

	c.mu.Lock()
	defer c.mu.Unlock()

	// 双重检查
	if db, exists := c.partitions[partitionKey]; exists {
		return db, nil
	}

	// 创建新的分区数据库
	dbPath := filepath.Join(c.dataDir, fmt.Sprintf("drill_state_%s.db", partitionKey))
	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open partition database %s: %v", partitionKey, err)
	}

	// 设置连接池参数
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)
	db.SetConnMaxLifetime(24 * time.Hour)

	// 创建表结构
	if err := c.createPartitionTable(db); err != nil {
		db.Close()
		return nil, fmt.Errorf("failed to create partition table: %v", err)
	}

	c.partitions[partitionKey] = db
	return db, nil
}

// createPartitionTable 创建分区表结构
func (c *DrillStateDbClient) createPartitionTable(db *sql.DB) error {
	createTableSQL := `
	CREATE TABLE IF NOT EXISTS drill_state_data (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		ct INTEGER NOT NULL,                    -- 时间戳（毫秒精度）
		dt TEXT NOT NULL,                       -- 数据传输时间（字符串 yyyymmddHHMMSS）
		mse_total REAL,                         -- 总MSE
		mse_downhole REAL,                      -- 下钻头的MSE
		d REAL,                                 -- 钻头直径，单位：in英寸
		r REAL,                                 -- 钻头半径，单位：in英寸
		tq_surf REAL,                           -- 地面扭矩，单位：ft-lb 英尺-磅
		rpm_surf REAL,                          -- 钻柱转速，单位：转/分钟
		kn REAL,                                -- 螺杆转速系数，单位：rev/gal，转/加仑
		ktt REAL,                               -- 螺杆扭矩系数，单位：ft-lb/psi
		delta_p REAL,                           -- 压力，单位：psi
		q REAL,                                 -- 流量，单位：加仑/分钟
		doc REAL,                               -- 切削深度，单位：in/转
		ds REAL,                                -- 钻井强度，单位：psi
		ssi REAL,                               -- 粘滑指数
		ucs REAL,                               -- 岩石抗压强度
		created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
		data_source TEXT
	);

	-- 创建索引
	CREATE INDEX IF NOT EXISTS idx_drill_state_ct ON drill_state_data(ct);
	CREATE INDEX IF NOT EXISTS idx_drill_state_dt ON drill_state_data(dt);
	CREATE INDEX IF NOT EXISTS idx_drill_state_created ON drill_state_data(created_at);
	`

	_, err := db.Exec(createTableSQL)
	return err
}

// getPartitionKey 根据时间获取分区键 (格式: YYYYMM)
func (c *DrillStateDbClient) getPartitionKey(t time.Time) string {
	return t.Format("200601")
}

// getPartitionKeysByTimeRange 获取时间范围内的所有分区键
func (c *DrillStateDbClient) getPartitionKeysByTimeRange(startTime, endTime time.Time) []string {
	partitionKeys := make(map[string]bool)

	current := time.Date(startTime.Year(), startTime.Month(), 1, 0, 0, 0, 0, startTime.Location())
	end := time.Date(endTime.Year(), endTime.Month(), 1, 0, 0, 0, 0, endTime.Location())

	for !current.After(end) {
		partitionKeys[c.getPartitionKey(current)] = true
		current = current.AddDate(0, 1, 0) // 下个月
	}

	keys := make([]string, 0, len(partitionKeys))
	for key := range partitionKeys {
		keys = append(keys, key)
	}

	return keys
}

// InsertDrillStateData 插入钻井状态数据
func (c *DrillStateDbClient) InsertDrillStateData(data *DrillStateData) error {
	partitionKey := c.getPartitionKey(time.Unix(data.Ct/1000, 0))
	db, err := c.getPartitionDB(partitionKey)
	if err != nil {
		return err
	}

	// 设置默认值
	if data.CreatedAt.IsZero() {
		data.CreatedAt = time.Now()
	}
	if data.DataSource == "" {
		data.DataSource = "drill_state"
	}

	query := `
	INSERT INTO drill_state_data (
		ct, dt, mse_total, mse_downhole, d, r, tq_surf, rpm_surf,
		kn, ktt, delta_p, q, doc, ds, ssi, ucs, created_at, data_source
	) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`

	_, err = db.Exec(query,
		data.Ct, data.Dt, data.MSETotal, data.MSEDownhole, data.D, data.R,
		data.TQsurf, data.RPMsurf, data.Kn, data.Kt, data.DeltaP,
		data.Q, data.DOC, data.DS, data.SSI, data.UCS, data.CreatedAt, data.DataSource,
	)

	if err != nil {
		return fmt.Errorf("failed to insert drill state data: %v", err)
	}

	return nil
}

// QueryByTimeRange 根据时间范围查询数据
func (c *DrillStateDbClient) QueryByTimeRange(startTime, endTime time.Time) ([]*DrillStateData, error) {
	partitionKeys := c.getPartitionKeysByTimeRange(startTime, endTime)

	var allResults []*DrillStateData

	for _, partitionKey := range partitionKeys {
		db, err := c.getPartitionDB(partitionKey)
		if err != nil {
			log.Printf("Warning: failed to get partition DB %s: %v", partitionKey, err)
			continue
		}

		query := `
		SELECT ct, dt, mse_total, mse_downhole, d, r, tq_surf, rpm_surf,
		       kn, ktt, delta_p, q, doc, ds, ssi, ucs, created_at, data_source
		FROM drill_state_data
		WHERE ct >= ? AND ct <= ?
		ORDER BY ct ASC
		`

		startMs := startTime.UnixMilli()
		endMs := endTime.UnixMilli()

		rows, err := db.Query(query, startMs, endMs)
		if err != nil {
			log.Printf("Warning: failed to query partition %s: %v", partitionKey, err)
			continue
		}
		defer rows.Close()

		for rows.Next() {
			var data DrillStateData
			err := rows.Scan(
				&data.Ct, &data.Dt, &data.MSETotal, &data.MSEDownhole, &data.D, &data.R,
				&data.TQsurf, &data.RPMsurf, &data.Kn, &data.Kt, &data.DeltaP,
				&data.Q, &data.DOC, &data.DS, &data.SSI, &data.UCS, &data.CreatedAt, &data.DataSource,
			)
			if err != nil {
				log.Printf("Warning: failed to scan row: %v", err)
				continue
			}
			allResults = append(allResults, &data)
		}
	}

	return allResults, nil
}

// QueryRecentSeconds 查询最近n秒的数据
func (c *DrillStateDbClient) QueryRecentSeconds(seconds int) ([]*DrillStateData, error) {
	endTime := time.Now()
	startTime := endTime.Add(-time.Duration(seconds) * time.Second)
	return c.QueryByTimeRange(startTime, endTime)
}

// QueryRecentMinutes 查询最近n分钟的数据
func (c *DrillStateDbClient) QueryRecentMinutes(minutes int) ([]*DrillStateData, error) {
	endTime := time.Now()
	startTime := endTime.Add(-time.Duration(minutes) * time.Minute)
	return c.QueryByTimeRange(startTime, endTime)
}

// QueryRecentDays 查询最近n天的数据
func (c *DrillStateDbClient) QueryRecentDays(days int) ([]*DrillStateData, error) {
	endTime := time.Now()
	startTime := endTime.AddDate(0, 0, -days)
	return c.QueryByTimeRange(startTime, endTime)
}

// GetPartitions 获取所有分区
func (c *DrillStateDbClient) GetPartitions() ([]string, error) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	partitions := make([]string, 0, len(c.partitions))
	for key := range c.partitions {
		partitions = append(partitions, key)
	}
	return partitions, nil
}

// GetPartitionStats 获取分区统计信息
func (c *DrillStateDbClient) GetPartitionStats(partitionKey string) (map[string]interface{}, error) {
	db, err := c.getPartitionDB(partitionKey)
	if err != nil {
		return nil, err
	}

	stats := make(map[string]interface{})

	// 获取记录数
	var count int
	err = db.QueryRow("SELECT COUNT(*) FROM drill_state_data").Scan(&count)
	if err != nil {
		return nil, err
	}
	stats["record_count"] = count

	// 获取时间范围
	var minTime, maxTime int64
	err = db.QueryRow("SELECT MIN(ct), MAX(ct) FROM drill_state_data").Scan(&minTime, &maxTime)
	if err != nil {
		return nil, err
	}

	if count > 0 {
		stats["min_time"] = time.Unix(minTime/1000, (minTime%1000)*1e6).Format("2006-01-02 15:04:05")
		stats["max_time"] = time.Unix(maxTime/1000, (maxTime%1000)*1e6).Format("2006-01-02 15:04:05")
	} else {
		stats["min_time"] = ""
		stats["max_time"] = ""
	}

	return stats, nil
}
